from openai import OpenAI
from dotenv import load_dotenv
import numpy as np
from numpy import dot
from numpy.linalg import norm

load_dotenv()

client = OpenAI()


def cos_sim(a, b):
    '''余弦距离 -- 越大越相似'''
    return dot(a, b) / (norm(a) * norm(b))


def l2(a, b):
    '''欧式距离 -- 越小越相似'''
    x = np.asarray(a) - np.asarray(b)
    return norm(x)


# 获取文本的向量表示
def get_embeddings(text, model="text-embedding-v3"):
    data = client.embeddings.create(input=text, model=model).data
    return [record.embedding for record in data]


query = "国际争端"

documents = [
    "联合国安理会上，俄罗斯与美国，伊朗与以色列“吵”起来了",
    "土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判",
    "日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤",
    "孙志刚被判死缓 减为无期徒刑后终身监禁 不得减刑、假释",
    "以色列立法禁联合国机构，美表态担忧，中东局势再生波澜",
]

# 获取查询的向量  第 0 位，表示文本向量的表示     因为 get_embedding 返回的是一个列表 [[]] 所以取第0个元素
query_vec = get_embeddings([query])[0]

doc_vecs = get_embeddings(documents)

print("Cosine distance:")
print(cos_sim(query_vec, query_vec))
for vec in doc_vecs:
    print(cos_sim(query_vec, vec))

